# -*- coding: utf-8 -*-
"""
Created on Mon Mar 22 13:44:05 2021

@author: ckpark
"""

import os
os.environ["PROJ_LIB"] = r'C:\Users\ckpark\anaconda3\Library\share'

from netCDF4 import Dataset
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import sys
import matplotlib.gridspec as gridspec
from scipy.stats import linregress
import matplotlib.colors as mc
from matplotlib.patches import Polygon
import matplotlib as m
import matplotlib
from scipy import stats
import scipy.signal
import scipy.signal as signal

calend = ['6/1/2021','1/1/2022','6/1/2022','12/1/2022','12/31/2022']
xlabs = ['JUN'+'\n'+'2021','JAN'+'\n'+'2022','JUN'+'\n'+'2022','DEC'+'\n'+'2022','DEC'+'\n'+'2022']
#calend = ['4/18/1942','4/18/1943','4/6/1944','5/9/1944']
#################################################################
def moving(inte, n):
    result = np.convolve(inte, np.ones((n,))/n, mode='same')
    return result

def passfuc(nt,dat):
    fs = nt
    # First, design the Buterworth filter
    N  = 5    # Filter order

    #cut = ((fs)/float(cutoff))
    cut = 10.0
    Wn = cut / (fs/2.0) # Cutoff frequency
    B, A = signal.butter(N, Wn, btype='low',output='ba')
    # Second, apply the filter
    rdat = signal.filtfilt(B,A, dat)
    return rdat


def init_plotting():
    plt.rcParams['font.family'] = 'Times New Roman'
init_plotting()
#################################################################
######################################
read = np.loadtxt('./RESULT/ori_EDI_47156', dtype='str')
yr = np.array(read[1:,0], dtype='int')
mo = np.array(read[1:,1], dtype='int')
dy = np.array(read[1:,2], dtype='int')
edi = np.array(read[1:,3], dtype='float')
pre = np.array(read[1:,4], dtype='float')

read = np.loadtxt('./RESULT/Clim_EDI_47156', dtype='str')
clim_edi = np.array(read[1:,3], dtype='float')
clim_pre = np.array(read[1:,4], dtype='float')

read = np.loadtxt('./RESULT/MAM_EDI_47156', dtype='str')
mam_edi = np.array(read[1:,3], dtype='float')
mam_pre = np.array(read[1:,4], dtype='float')

read = np.loadtxt('./RESULT/JJA_EDI_47156', dtype='str')
jja_edi = np.array(read[1:,3], dtype='float')
jja_pre = np.array(read[1:,4], dtype='float')

read = np.loadtxt('./RESULT/MAMJJA_EDI_47156', dtype='str')
mamjja_edi = np.array(read[1:,3], dtype='float')
mamjja_pre = np.array(read[1:,4], dtype='float')

dates = []

vvar = []
for t in range(len(yr)):
    for k in range(len(calend)):
        var = calend[k].split('/')
        year = int(var[-1]); month = int(var[0]); day = int(var[1])
        if (yr[t] == year) and (mo[t] == month) and (dy[t] == day):
            dates.append(t) 
            
        if yr[t] == 2022 and mo[t] == 3 and dy[t] == 1:
            vvar.append(t)

        if yr[t] == 2022 and mo[t] == 6 and dy[t] == 1:
            vvar.append(t)            

        if yr[t] == 2022 and mo[t] == 8 and dy[t] == 31:
            vvar.append(t)

edi = edi[dates[0]:dates[-1]+1]
pre = pre[dates[0]:dates[-1]+1]

clim_edi = clim_edi[dates[0]:dates[-1]+1]
mam_edi = mam_edi[dates[0]:dates[-1]+1]
jja_edi = jja_edi[dates[0]:dates[-1]+1]
mamjja_edi = mamjja_edi[dates[0]:dates[-1]+1]

clim_pre = clim_pre[dates[0]:dates[-1]+1]
mam_pre = mam_pre[dates[0]:dates[-1]+1]
jja_pre = jja_pre[dates[0]:dates[-1]+1]
mamjja_pre = mamjja_pre[dates[0]:dates[-1]+1]

nt = dates[-1] - dates[0] + 1

acc_pre = np.zeros((len(pre)))
acc_clim = np.zeros((len(pre)))
acc_mam = np.zeros((len(pre)))
acc_jja = np.zeros((len(pre)))

acc_pre[0] = pre[0]
acc_clim[0] = clim_pre[0]
acc_mam[0] = mam_pre[0]
acc_jja[0] = jja_pre[0]
for t in range(len(pre)-1):
    acc_pre[t+1] = acc_pre[t] + pre[t+1]
    acc_clim[t+1] = acc_clim[t] + clim_pre[t+1]
    acc_mam[t+1] = acc_mam[t] + mam_pre[t+1]
    acc_jja[t+1] = acc_jja[t] + jja_pre[t+1]


############## PDO FIG #############
fig = plt.figure(figsize=(10,10))
for xx in range(1):
    ax1 = fig.add_subplot(3,1,xx+1)

    xax = np.arange(0,nt)
    ax1.set_xlim(0,nt-1)
    ax1.set_yticks(np.arange(-10.0,10.0,1.0))
    ax1.set_ylim(-2.1,2.1)
  
    xlab = calend[:-1]
    xtick = np.array(dates[:-1]) - dates[0]

    ax1.set_xticks(xtick)
    ax1.set_xticklabels(xlabs)
    ax1.set_xlabel('Date',fontsize=26)    
    #plt.xticks(rotation=45)
    
    ax1.axhline(y=0.0, color='black', linestyle='solid')
    ax1.axhline(y=-1.0, color='red', linestyle='dashed')
    ax1.axhline(y=-2.0, color='red', linestyle='dashed')
    
    for k in range(len(vvar)):
        ax1.axvline(x=vvar[k]-dates[0], color='gray', linestyle='solid')
        

    ax1.plot(xax,edi,color='black',zorder=2,alpha=1.0,label='Original')
    
    ax1.plot(xax,clim_edi,color='cyan',zorder=1,alpha=1.0,label='Clim(all month)')
    ax1.plot(xax,mam_edi,color='green',zorder=1,alpha=1.0,label='Clim(MAM)')
    ax1.plot(xax,jja_edi,color='blue',zorder=1,alpha=1.0,label='Clim(JJA)')
    #ax1.plot(xax,mamjja_edi,color='orange',zorder=1,alpha=1.0,label='Clim(MAMJJA)')
    #ax1.plot(xax,macd_rpdo,color='blue',zorder=1,alpha=1.0)
    ax1.set_ylabel('Index',fontsize=24)
    ax1.tick_params(labelsize=23)

    #ax1.axvline(x=xtick[1], color='black', linestyle='solid')
    #ax1.axvline(x=xtick[2], color='black', linestyle='solid')
    #ax1.axvline(x=276, color='black', linestyle='solid')
    
    #plt.legend(loc='lower left', bbox_to_anchor=(1, 0.3),fontsize=15)
    '''
    ######
    ax2 = ax1.twinx()
    
    ax2.set_yticks(np.arange(0,500,50))
    ax2.set_ylim(0,200)
    
    ax2.bar(xax,pre,width=1.5,color='green',zorder=1,alpha=1)
    ax2.set_ylabel('Precipitation'+'\n'+'(mm)',fontsize=24)
    ax2.tick_params(labelsize=25)
    '''

plt.subplots_adjust(hspace=0.0,wspace=0.2)
plt.show()
outfig = '../99.FIG/Clim_SCS_past_EDI_time_series.png'      
fig.savefig(outfig,bbox_inches='tight') 
#############################################################33

############## PDO FIG #############
fig = plt.figure(figsize=(10,10))
for xx in range(1):
    ax1 = fig.add_subplot(3,1,xx+1)

    xax = np.arange(0,nt)
    ax1.set_xlim(0,nt-1)
    ax1.set_yticks(np.arange(0,2300,500))
    ax1.set_ylim(0.1,2400)
  
    xlab = calend[:-1]
    xtick = np.array(dates[:-1]) - dates[0]

    ax1.set_xticks(xtick)
    ax1.set_xticklabels(xlabs)
    ax1.set_xlabel('Date',fontsize=26)    
    #plt.xticks(rotation=45)
    
    #ax1.axhline(y=0.0, color='black', linestyle='solid')
    #ax1.axhline(y=-1.0, color='red', linestyle='dashed')
    #ax1.axhline(y=-2.0, color='red', linestyle='dashed')
    
    for k in range(len(vvar)):
        ax1.axvline(x=vvar[k]-dates[0], color='gray', linestyle='solid')
        

    #ax1.bar(xax,pre,color='black',zorder=3,alpha=1.0,label='Original')
    #ax1.bar(xax,clim_pre,color='cyan',zorder=1,alpha=1.0,label='Clim(all month)')
    #ax1.bar(xax,mam_pre,color='green',zorder=2,alpha=1.0,label='Clim(MAM)')
    #ax1.bar(xax,jja_pre,color='blue',zorder=2,alpha=1.0,label='Clim(JJA)')
    
    ax1.plot(xax,acc_pre,color='black',zorder=3,alpha=1.0,label='Original')
    
    ax1.plot(xax,acc_clim,color='cyan',zorder=1,alpha=1.0,label='Clim(all month)')
    ax1.plot(xax,acc_mam,color='green',zorder=2,alpha=1.0,label='Clim(MAM)')
    ax1.plot(xax,acc_jja,color='blue',zorder=2,alpha=1.0,label='Clim(JJA)')
    
    #ax1.bar(xax,mamjja_pre,color='orange',zorder=1,alpha=1.0,label='Clim(MAMJJA)')
    #ax1.plot(xax,macd_rpdo,color='blue',zorder=1,alpha=1.0)
    ax1.set_ylabel('Accumulated'+'\n'+'precipitation (mm)',fontsize=24)
    ax1.tick_params(labelsize=23)
    
    #ax1.axvline(x=xtick[1], color='black', linestyle='solid')
    #ax1.axvline(x=xtick[2], color='black', linestyle='solid')
    #ax1.axvline(x=276, color='black', linestyle='solid')
    
    #plt.legend(loc='lower left', bbox_to_anchor=(1, 0.3),fontsize=15)
    '''
    ######
    ax2 = ax1.twinx()
    
    ax2.set_yticks(np.arange(0,500,50))
    ax2.set_ylim(0,200)
    
    ax2.bar(xax,pre,width=1.5,color='green',zorder=1,alpha=1)
    ax2.set_ylabel('Precipitation'+'\n'+'(mm)',fontsize=24)
    ax2.tick_params(labelsize=25)
    '''

plt.subplots_adjust(hspace=0.0,wspace=0.2)
plt.show()
outfig = '../99.FIG/Clim_SCS_past_PRE_time_series.png'      
fig.savefig(outfig,bbox_inches='tight') 
#############################################################33

